import os
import torch
import numpy as np
import open3d as o3d
from torch.utils.data import Dataset, DataLoader
from tqdm import tqdm

# 加载 ModelNet10 数据集
class ModelNet10(Dataset):
    def __init__(self, root_dir, split='train', num_points=1024):
        """
        初始化数据集
        :param root_dir: 数据集根目录
        :param split: 'train' 或 'test'
        :param num_points: 每个点云的点数
        """
        self.root_dir = root_dir
        self.split = split
        self.num_points = num_points
        
        # 数据存储
        self.data = []
        self.labels = []
        
        # 获取类别名
        self.classes = sorted(os.listdir(root_dir))
        
        # 读取每个类别的训练或测试数据
        for label, cls in enumerate(self.classes):
            cls_dir = os.path.join(root_dir, cls)
            split_dir = os.path.join(cls_dir, split)  # 选择训练集或测试集
            if os.path.exists(split_dir):
                for file in os.listdir(split_dir):
                    if file.endswith('.off'):
                        # 添加文件路径和标签
                        self.data.append(os.path.join(split_dir, file))
                        self.labels.append(label)
        
    def __len__(self):
        return len(self.data)
    
    def __getitem__(self, idx):
        # 获取文件路径和标签
        file_path = self.data[idx]
        label = self.labels[idx]
        
        # 使用 Open3D 加载 .off 文件
        mesh = o3d.io.read_triangle_mesh(file_path)
        pcd = mesh.sample_points_poisson_disk(number_of_points=self.num_points)
        
        # 转换为 numpy 数组
        points = np.asarray(pcd.points)
        
        # 归一化处理
        points = points - np.mean(points, axis=0)  # 中心化
        points = points / np.max(np.linalg.norm(points, axis=1))  # 归一化到单位球体
        
        # 转换为 Tensor 格式
        points = torch.FloatTensor(points.T)  # [3, num_points]
        label = torch.tensor(label, dtype=torch.long)
        
        return points, label

# 数据加载
def load_data(root_dir, split='train', batch_size=32, num_points=1024):
    """
    返回数据加载器
    :param root_dir: 数据集路径
    :param split: 'train' 或 'test'
    :param batch_size: 批量大小
    :param num_points: 每个点云的点数
    :return: DataLoader
    """
    dataset = ModelNet10(root_dir=root_dir, split=split, num_points=num_points)
    dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=True)
    return dataloader

if __name__ == "__main__": 
    dataloader = load_data("C:\\Users\\maohai_pang\\Desktop\\研二(上)\\项目\\github\\datasets\\ModelNet10")
    print("成功加载.......")
    pass
